"""
AutoGen 工具函数模块
包含常用的 AutoGen 辅助函数和配置
"""

import autogen
import os
from typing import List, Dict, Any

class AutoGenManager:
    """AutoGen 管理器类"""
    
    def __init__(self, config: Dict[str, Any] = None):
        """初始化 AutoGen 管理器"""
        self.config = config or self._default_config()
        self.agents = {}
    
    def _default_config(self) -> Dict[str, Any]:
        """默认配置"""
        return {
            "model": os.getenv("AUTOGEN_MODEL", "gpt-4"),
            "api_key": os.getenv("OPENAI_API_KEY"),
            "temperature": float(os.getenv("AUTOGEN_TEMPERATURE", "0.7")),
            "max_tokens": int(os.getenv("AUTOGEN_MAX_TOKENS", "1000")),
        }
    
    def create_llm_config(self) -> Dict[str, Any]:
        """创建 LLM 配置"""
        config_list = [
            {
                "model": self.config["model"],
                "api_key": self.config["api_key"],
            }
        ]
        
        return {
            "config_list": config_list,
            "temperature": self.config["temperature"],
            "max_tokens": self.config["max_tokens"],
        }
    
    def create_user_proxy(self, name: str = "user_proxy", **kwargs) -> autogen.UserProxyAgent:
        """创建用户代理"""
        default_config = {
            "name": name,
            "system_message": "A human admin.",
            "code_execution_config": {
                "last_n_messages": 2,
                "work_dir": "groupchat",
                "use_docker": False,
            },
            "human_input_mode": "NEVER",
        }
        default_config.update(kwargs)
        
        agent = autogen.UserProxyAgent(**default_config)
        self.agents[name] = agent
        return agent
    
    def create_assistant(self, name: str = "assistant", **kwargs) -> autogen.AssistantAgent:
        """创建助手代理"""
        default_config = {
            "name": name,
            "llm_config": self.create_llm_config(),
        }
        default_config.update(kwargs)
        
        agent = autogen.AssistantAgent(**default_config)
        self.agents[name] = agent
        return agent
    
    def create_group_chat(self, agents: List[autogen.Agent], **kwargs) -> autogen.GroupChat:
        """创建群聊"""
        default_config = {
            "agents": agents,
            "messages": [],
            "max_round": 10,
        }
        default_config.update(kwargs)
        
        return autogen.GroupChat(**default_config)
    
    def get_agent(self, name: str) -> autogen.Agent:
        """获取已创建的智能体"""
        return self.agents.get(name)

def create_simple_chat():
    """创建简单的对话示例"""
    manager = AutoGenManager()
    
    # 创建智能体
    user_proxy = manager.create_user_proxy()
    assistant = manager.create_assistant()
    
    return user_proxy, assistant

def create_multi_agent_system():
    """创建多智能体系统示例"""
    manager = AutoGenManager()
    
    # 创建多个智能体
    user_proxy = manager.create_user_proxy()
    coder = manager.create_assistant(
        name="coder",
        system_message="You are a helpful coding assistant."
    )
    reviewer = manager.create_assistant(
        name="reviewer", 
        system_message="You are a code reviewer."
    )
    
    # 创建群聊
    group_chat = manager.create_group_chat([user_proxy, coder, reviewer])
    manager_agent = autogen.GroupChatManager(
        groupchat=group_chat,
        llm_config=manager.create_llm_config()
    )
    
    return user_proxy, manager_agent
